import zmq
import json
import time
from threading import Thread
from dataclasses import dataclass

@dataclass
class MarketData:
    symbol: str
    price: float
    volume: float
    timestamp: int
    exchange: str
    data_type: str  # 'tick', 'bar', 'orderbook'
    additional_fields: dict = None

class MarketDataFeed:
    """市场数据接入基类"""
    def __init__(self, config):
        self.config = config
        self.subscribers = []
        self.is_running = False
        self.context = zmq.Context()
        self.socket = None
        self.connection_status = False
        
    def connect(self):
        """连接到数据源"""
        try:
            # 根据配置创建对应的socket连接
            if self.config['protocol'] == 'zmq':
                self.socket = self.context.socket(zmq.SUB)
                self.socket.connect(self.config['endpoint'])
                for symbol in self.config['symbols']:
                    self.socket.setsockopt_string(zmq.SUBSCRIBE, symbol)
            elif self.config['protocol'] == 'kafka':
                # Kafka连接实现
                pass
            
            self.connection_status = True
            print(f"Connected to market data source: {self.config['name']}")
            return True
        except Exception as e:
            print(f"Failed to connect to market data source: {str(e)}")
            return False
    
    def start(self):
        """启动数据接收线程"""
        if not self.is_running and self.connection_status:
            self.is_running = True
            self.receive_thread = Thread(target=self._receive_data)
            self.receive_thread.daemon = True
            self.receive_thread.start()
            print(f"Market data feed started: {self.config['name']}")
    
    def stop(self):
        """停止数据接收"""
        self.is_running = False
        if hasattr(self, 'receive_thread'):
            self.receive_thread.join(timeout=2.0)
        if self.socket:
            self.socket.close()
        print(f"Market data feed stopped: {self.config['name']}")
    
    def _receive_data(self):
        """数据接收循环"""
        while self.is_running:
            try:
                # 接收数据并解析
                data = self._receive_raw_data()
                if data:
                    market_data = self._parse_data(data)
                    self._notify_subscribers(market_data)
            except Exception as e:
                print(f"Error receiving market data: {str(e)}")
                time.sleep(0.001)  # 短暂休眠避免CPU占用过高
    
    def _receive_raw_data(self):
        """接收原始数据，子类实现"""
        raise NotImplementedError("子类必须实现_receive_raw_data方法")
    
    def _parse_data(self, raw_data):
        """解析原始数据，子类实现"""
        raise NotImplementedError("子类必须实现_parse_data方法")
    
    def subscribe(self, subscriber):
        """添加数据订阅者"""
        if subscriber not in self.subscribers:
            self.subscribers.append(subscriber)
    
    def unsubscribe(self, subscriber):
        """移除数据订阅者"""
        if subscriber in self.subscribers:
            self.subscribers.remove(subscriber)
    
    def _notify_subscribers(self, market_data):
        """通知所有订阅者"""
        for subscriber in self.subscribers:
            try:
                subscriber.on_market_data(market_data)
            except Exception as e:
                print(f"Error notifying subscriber: {str(e)}")
    
    def get_connection_status(self):
        """获取连接状态"""
        return self.connection_status